package org.apache.submarine.spark.security

import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.optimizer.{SubmarineConfigurationCheckExtension, SubmarineDataMaskingExtension, SubmarinePushPredicatesThroughExtensions, SubmarineRowFilterExtension, SubmarineSparkRangerAuthorizationExtension}
import org.apache.spark.sql.execution.SubmarineSparkPlanOmitStrategy

/**
 * Spark SQL测试Extension扩展
 */
object ExtensionTest {
    def main(args: Array[String]): Unit = {
        // 设置执行用户
        // 需要在Ranger中添加好test用户
        System.setProperty("HADOOP_USER_NAME", "test")
        // 加载默认配置
        val conf = new SparkConf()

        // 创建Spark Session
        val spark = SparkSession.builder()
            .master("local[3]")
            .config(conf)
            .appName("Spark SQL Ranger测试")
            .withExtensions { ext =>
                ext.injectCheckRule(SubmarineConfigurationCheckExtension)
                ext.injectOptimizerRule(SubmarineSparkRangerAuthorizationExtension)
                ext.injectOptimizerRule(SubmarineRowFilterExtension)
                ext.injectOptimizerRule(SubmarineDataMaskingExtension)
                ext.injectOptimizerRule(SubmarinePushPredicatesThroughExtensions)
                ext.injectPlannerStrategy(SubmarineSparkPlanOmitStrategy)
            }
            .enableHiveSupport()
            .getOrCreate()

        spark.sql("insert overwrite rtdw.ddl_table_test values ('1'),('2');").show()
        // spark.sql("select * from rtdw.ddl_table_test;").show()
    }
}
